package com.hhh.springai_test.controller;

import com.hhh.springai_test.config.ChatConfig;
import com.hhh.springai_test.model.dto.ChatRequest;
import jakarta.annotation.Resource;
import okhttp3.*;
import okio.BufferedSink;
import okio.BufferedSource;
import org.jetbrains.annotations.NotNull;
import org.json.JSONObject;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@RestController
@RequestMapping("/ai")
public class deepseekController {

    // 用于存储会话历史的Map，key为sessionId，value为消息列表
    private static final Map<String, List<ChatMessage>> chatHistory = new ConcurrentHashMap<>();
    // 最大历史消息数量
    private static final int MAX_HISTORY_SIZE = 30;

    @Resource
    private ChatConfig chatConfig;

    // 新增ChatMessage类用于存储单条消息
    public static class ChatMessage {
        private String content;
        private String role;

        // getters and setters
        public String getContent() { return content; }
        public void setContent(String content) { this.content = content; }
        public String getRole() { return role; }
        public void setRole(String role) { this.role = role; }
    }

    // 修改接口方法，接收新的请求格式
    @PostMapping(value = "/deepSeekChat", produces = org.springframework.http.MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> generateStreamAsString2(@org.springframework.web.bind.annotation.RequestBody ChatRequest chatRequest) throws Exception {
        final String sessionId = chatRequest.getSessionId();
        List<ChatMessage> newMessages = chatRequest.getMessages();
        
        // 获取或创建会话历史
        List<ChatMessage> sessionHistory = chatHistory.computeIfAbsent(sessionId, k -> new ArrayList<>());
        
        // 添加所有新消息到历史记录
        sessionHistory.addAll(newMessages);
        
        // 如果历史记录超过最大限制，删除最早的消息
        while (sessionHistory.size() > MAX_HISTORY_SIZE) {
            sessionHistory.remove(0);
        }

        // 构建完整的消息历史JSON
        StringBuilder messagesJson = new StringBuilder();
        messagesJson.append("  \"messages\": [\n");
        
        // 添加历史消息
        for (int i = 0; i < sessionHistory.size(); i++) {
            ChatMessage msg = sessionHistory.get(i);
            messagesJson.append("    {\n")
                    .append("      \"content\": \"").append(escapeJsonString(msg.getContent())).append("\",\n")
                    .append("      \"role\": \"").append(msg.getRole()).append("\"\n")
                    .append("    }");
            if (i < sessionHistory.size() - 1) {
                messagesJson.append(",");
            }
            messagesJson.append("\n");
        }
        messagesJson.append("  ],\n");

        // 修改请求体构建部分
        OkHttpClient client = new OkHttpClient().newBuilder().build();
        RequestBody body = new RequestBody() {
            @Override
            public okhttp3.MediaType contentType() {
                return okhttp3.MediaType.parse("application/json");
            }

            @Override
            public void writeTo(BufferedSink sink) throws IOException {
                String requestBodyJson = "{\n" +
                        messagesJson.toString() +
                        "  \"model\": \"" +
                        chatConfig.getDeepSeekConfigModel() +
                        "\",\n" +
                        "  \"frequency_penalty\": 0,\n" +
                        "  \"max_tokens\": 1024,\n" +
                        "  \"presence_penalty\": 0,\n" +
                        "  \"response_format\": {\n" +
                        "    \"type\": \"text\"\n" +
                        "  },\n" +
                        "  \"stop\": null,\n" +
                        "  \"stream\": true,\n" +
                        "  \"stream_options\": null,\n" +
                        "  \"temperature\": 1,\n" +
                        "  \"top_p\": 1,\n" +
                        "  \"tools\": null,\n" +
                        "  \"tool_choice\": \"none\",\n" +
                        "  \"logprobs\": false,\n" +
                        "  \"top_logprobs\": null\n" +
                        "}";
                sink.writeUtf8(requestBodyJson);
            }
        };
        // 创建 Headers
        Headers headers = new Headers.Builder()
                .add("Content-Type", "application/json")
                .add("Accept", "application/json")
                .add("Authorization", "Bearer " + chatConfig.getDeepSeekConfigUrl())  // 使用您的 API 密钥
                .build();
        // 创建 HttpUrl
        HttpUrl url = HttpUrl.parse(chatConfig.getDeepSeekConfigHost()); // 使用您的 API URL
        if (url == null) {
            throw new IllegalArgumentException("您此时未携带URL参数");
        }
        // 构造 Request
        okhttp3.Request request = new okhttp3.Request.Builder()
                .url(url)
                .post(body)  // 使用流式请求体
                .headers(headers)
                .build();
        // 执行请求并返回 Flux 流
        return Flux.create(sink -> {
            client.newCall(request).enqueue(new Callback() {
                @Override
                public void onFailure(@NotNull Call call, @NotNull IOException e) {
                    // 异常发生时调用 sink.error()
                    sink.error(e);
                }
                @Override
                public void onResponse(@NotNull Call call, @NotNull Response response) throws IOException {
                    ResponseBody responseBody = response.body();
                    BufferedSource source = responseBody.source();
                    StringBuilder fullResponse = new StringBuilder(); // 用于收集完整的AI响应
                    
                    while (true) {
                        // 逐行读取响应数据
                        String line = source.readUtf8Line();
                        if (line == null) {
                            break;
                        } else if (line.startsWith("data:")) {
                            String data = line.substring(5).trim();
                            // 检查数据是否为结束信号"[DONE]"
                            if (!"[DONE]".equals(data)) {
                                // 解析接收到的数据为JSON对象
                                JSONObject jsonResponse = new JSONObject(data);
                                String finishReason = jsonResponse.getJSONArray("choices").getJSONObject(0).optString("finish_reason");
                                // 收集AI的回答内容
                                String content = jsonResponse.getJSONArray("choices")
                                    .getJSONObject(0)
                                    .getJSONObject("delta")
                                    .optString("content", "");
                                fullResponse.append(content);
                                if ("stop".equals(finishReason)) {
                                    // 结束时，将AI的完整回答添加到聊天历史
                                    ChatMessage aiMessage = new ChatMessage();
                                    aiMessage.setRole("assistant");
                                    aiMessage.setContent(fullResponse.toString());
                                    // 获取当前会话的历史记录并添加AI的回答
                                    List<ChatMessage> sessionHistory = chatHistory.get(sessionId);
                                    if (sessionHistory != null) {
                                        sessionHistory.add(aiMessage);
                                        // 如果超过最大限制，删除最早的消息
                                        while (sessionHistory.size() > MAX_HISTORY_SIZE) {
                                            sessionHistory.remove(0);
                                        }
                                    }
                                    // 结束时推送数据并完成
                                    sink.next(data);
                                    sink.complete();
                                    break;
                                } else {
                                    // 否则继续推送数据
                                    sink.next(data);
                                }
                                // 打印调试信息
                                System.out.println(jsonResponse.getJSONArray("choices"));
                            }
                        }
                    }
                }
            });
        });
    }

    // 添加JSON字符串转义方法
    private String escapeJsonString(String input) {
        if (input == null) {
            return "";
        }
        return input.replace("\\", "\\\\")
                   .replace("\"", "\\\"")
                   .replace("\n", "\\n")
                   .replace("\r", "\\r")
                   .replace("\t", "\\t");
    }

    @Scheduled(cron = "0 */10 * * * ?")
    public void clearChatHistory() {
        chatHistory.clear();
        System.out.println("已清理聊天历史记录");
    }

}
